import os
import time
import threading
import cv2
import numpy as np
from PIL import ImageGrab



class SeleniumVideoRecorder:
    def __init__(self, driver, output_file="recording.mp4", fps=15.0, resolution=(1920, 1080)):
        """
        初始化视频录制工具

        :param driver: Selenium WebDriver 实例
        :param output_file: 输出视频文件路径
        :param fps: 视频帧率
        :param resolution: 视频分辨率 (宽, 高)
        """
        self.driver = driver
        self.output_file = output_file
        self.fps = fps
        self.resolution = resolution
        self.is_recording = False
        self.thread = None
        self.frames = []

        # 修改driver的quit方法
        self._patch_driver_quit()

        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_file) if os.path.dirname(output_file) else None)

    def start(self):
        """开始录制"""
        if self.is_recording:
            return

        self.is_recording = True
        self.thread = threading.Thread(target=self._record_screen)
        self.thread.start()
        print(f"开始录制视频，保存到: {self.output_file}")

    def stop(self):
        """停止录制并保存视频"""
        if not self.is_recording:
            return

        self.is_recording = False
        if self.thread and self.thread.is_alive():
            self.thread.join()

        self._save_video()
        print(f"视频录制完成，已保存到: {self.output_file}")

    def _record_screen(self):
        """屏幕录制线程"""
        while self.is_recording:
            try:
                # 检查driver是否有效
                if not hasattr(self.driver, 'session_id'):
                    self.stop()
                    break

                # 捕获屏幕
                img = ImageGrab.grab()
                frame = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

                # 调整分辨率
                frame = cv2.resize(frame, self.resolution)

                # 添加到帧列表
                self.frames.append(frame)

                # 控制帧率
                time.sleep(1.0 / self.fps)
            except Exception as e:
                print(f"录制错误: {e}")
                self.stop()
                break

    def _save_video(self):
        """保存视频文件"""
        if not self.frames:
            return

        # 创建视频写入器
        fourcc = cv2.CV_64F  # 或使用 'avc1'
        out = cv2.VideoWriter(self.output_file, fourcc, self.fps, self.resolution)

        # 写入所有帧
        for frame in self.frames:
            out.write(frame)

        # 释放资源
        out.release()
        self.frames = []

    def _patch_driver_quit(self):
        """修改driver的quit方法以自动停止录制"""
        original_quit = self.driver.quit

        def patched_quit():
            self.stop()
            original_quit()
            print("浏览器已退出，视频录制停止")

        self.driver.quit = patched_quit


